/*
 * NETCAP - Traffic Analysis Framework
 * Copyright (c) 2017 Philipp Mieden <dreadl0ck [at] protonmail [dot] ch>
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package label

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/dreadl0ck/netcap/utils"
	"github.com/evilsocket/islazy/tui"
	"github.com/pkg/errors"
	pb "gopkg.in/cheggaaa/pb.v1"
)

// eg: 04/15/2014-11:53:20.462091
const suricataTS = "01/02/2006-15:04:05.000000"

// regular expressions to match data from suricata fast.log
var (
	protoc         = regexp.MustCompile(`\{[A-Z]*\}`)
	classification = regexp.MustCompile(`\[Classification:[\s\w]*\]`)
	flowIdent      = regexp.MustCompile(`\}.*`)
	description    = regexp.MustCompile(`\[\*\*\].*\[\*\*\]`)

	// in case more than one label for the same timestamp exists
	// stop execution and print info
	// this affects layers being labeled, because they use the labelMap
	// other record types use the label array, which is not affected.
	// handling this needs to be improved in the future
	StopOnDuplicateLabels = false

	DisableLayerMapping = false
)

// SuricataAlert is a summary structure of an alerts contents
type SuricataAlert struct {
	Timestamp      string
	Proto          string
	SrcIP          string
	SrcPort        int
	DstIP          string
	DstPort        int
	Classification string
	Description    string
}

// Suricata creates labeled CSV files for audit records derived from the provided input file
// alerts are generated by using suricata to scan the input pcap file
// a directory named after the input file is created, all suricata logs go there
// if no output directory is specified, netcap audit records are expected in the current directory.
// otherwise audit records are expected in the output directory
func Suricata(inputPcap string, outputPath string, useDescription bool, separator, selection string) error {
	start := time.Now()

	// directory for suricata logs
	logDir := strings.TrimSuffix(strings.TrimSuffix(inputPcap, ".pcap"), ".pcapng")

	debug("checking log dir:", logDir)
	if _, err := os.Stat(logDir); err != nil {
		// does not exist
		// make sure suricata log dir exists
		if os.IsNotExist(err) {
			if err := os.Mkdir(logDir, 0755); err != nil {
				return err
			}
		} else {
			return err
		}

	} else {
		debug("removing suricata logfiles from previous runs")

		// clean suricata logfiles from previous runs
		// errors are explicitely ignored, because if the files do not exist
		// there will be an error
		_ = os.Remove(logDir + "/fast.log")
		_ = os.Remove(logDir + "/eve.json")
		_ = os.Remove(logDir + "/stats.log")
	}

	fmt.Println("scanning", inputPcap, "with suricata...")

	// call suricata to scan dump file
	out, err := exec.Command("suricata", "-c", "/usr/local/etc/suricata/suricata.yaml", "-r", inputPcap, "-l", logDir).CombinedOutput()
	if err != nil {
		return errors.Wrap(err, "Error with output: "+string(out))
	}

	// get path for logfile
	path := filepath.Join(logDir, "fast.log")
	fmt.Println("done. reading logs from", path)

	// TODO: switch to eve logs in the future
	// pathEve := filepath.Join(logDir, "eve.json")
	// eveChan, errChan := surevego.LoadEveJSONFile(pathEve)
	// for e := range eveChan {
	// 	fmt.Println(e)
	// }
	// for e := range errChan {
	// 	fmt.Println("error:", e)
	// }

	// os.Exit(0)

	// read fast.log contents
	contents, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}

	// extract alerts
	labelMap, labels, err := ParseSuricataFastLog(contents, useDescription)
	if err != nil {
		return err
	}

	if len(labels) == 0 {
		fmt.Println("no labels found.")
		os.Exit(0)
	}

	fmt.Println("got", len(labels), "labels")

	rows := [][]string{}
	for c, num := range ClassificationMap {
		rows = append(rows, []string{c, strconv.Itoa(num)})
	}

	// print alert summary
	tui.Table(os.Stdout, []string{"Classification", "Count"}, rows)
	fmt.Println()

	// apply labels to data
	// set outDir to current dir or flagOut
	var outDir string
	if outputPath != "" {
		outDir = outputPath
	} else {
		outDir = "."
	}

	// label all layer data in outDir
	// first read directory
	files, err := ioutil.ReadDir(outDir)
	if err != nil {
		return err
	}

	var (
		wg  sync.WaitGroup
		pbs []*pb.ProgressBar
	)

	// iterate over all files in dir
	for _, f := range files {
		// check if its an audit record file
		if strings.HasSuffix(f.Name(), ".ncap.gz") || strings.HasSuffix(f.Name(), ".ncap") {
			wg.Add(1)

			var (
				// get record name
				filename = f.Name()
				typ      = strings.TrimSuffix(strings.TrimSuffix(filename, ".ncap.gz"), ".ncap")
			)

			// some record types need to be processed separately
			// because the mapping logic differs for them
			switch typ {
			case "UDP":
				pbs = append(pbs, UDP(&wg, filename, labels, outputPath, separator, selection))
			case "TCP":
				pbs = append(pbs, TCP(&wg, filename, labels, outputPath, separator, selection))
			case "IPv4":
				pbs = append(pbs, IPv4(&wg, filename, labels, outputPath, separator, selection))
			case "IPv6":
				pbs = append(pbs, IPv6(&wg, filename, labels, outputPath, separator, selection))
			case "Connection":
				pbs = append(pbs, Connections(&wg, filename, labels, outputPath, separator, selection))
			case "Flow":
				pbs = append(pbs, Flows(&wg, filename, labels, outputPath, separator, selection))
			case "HTTP":
				pbs = append(pbs, HTTP(&wg, filename, labels, outputPath, separator, selection))
			case "TLS":
				pbs = append(pbs, TLS(&wg, filename, labels, outputPath, separator, selection))
			// LinkFlows can currently not be labeled with suricata because the alerts dont have L2 information
			// case "LinkFlow":
			// 	pbs = append(pbs, LinkFlow(&wg, filename, labels, outputPath, separator, selection))
			case "NetworkFlow":
				pbs = append(pbs, NetworkFlow(&wg, filename, labels, outputPath, separator, selection))
			case "TransportFlow":
				pbs = append(pbs, TransportFlow(&wg, filename, labels, outputPath, separator, selection))
			default:
				if !DisableLayerMapping {
					// apply labels to all records by timestamp only
					pbs = append(pbs, Layer(&wg, filename, typ, labelMap, labels, outputPath, separator, selection))
				}
			}
		}
	}

	var pool *pb.Pool
	if UseProgressBars {

		// wait for goroutines to start and initialize
		// otherwise progress bars will bug
		time.Sleep(3 * time.Second)

		// start pool
		pool, err = pb.StartPool(pbs...)
		if err != nil {
			return err
		}
		utils.ClearScreen()
	}

	wg.Wait()

	if UseProgressBars {
		// close pool
		if err := pool.Stop(); err != nil {
			fmt.Println("failed to stop progress bar pool:", err)
		}
	}

	fmt.Println("\ndone in", time.Since(start))
	return nil
}

// ParseSuricataFastLog returns labels for a given suricata fast.log contents.
func ParseSuricataFastLog(contents []byte, useDescription bool) (labelMap map[string]*SuricataAlert, arr []*SuricataAlert, err error) {
	fmt.Println("parsing suricata fast.log")

	if len(excluded) != 0 {
		var excludedNames []string
		for n := range excluded {
			excludedNames = append(excludedNames, n)
		}
		fmt.Println("excluding alerts with the following classifications:", excludedNames)
	}

	// alerts that have a duplicate timestamp
	var duplicates = []*SuricataAlert{}

	// ts:alert
	labelMap = make(map[string]*SuricataAlert)

	// range fast.log contents line by line
	for _, l := range strings.Split(string(contents), "\n") {

		// minimum 27 chars for a valid log line starting with a timestamp
		if len(l) > 27 {

			// parse timestamp
			// important: parse in current location
			// declaring t here to avoid shadowing the error
			var t = time.Time{}
			t, err = time.ParseInLocation(suricataTS, l[:26], time.Local)
			if err != nil {
				return
			}

			// extract protocol name
			nProto := strings.TrimSuffix(strings.TrimPrefix(protoc.FindString(l), "{"), "}")

			// extract flow identifier
			flow := strings.TrimPrefix(flowIdent.FindString(l), "} ")

			// extract values
			flowSlice := strings.Split(flow, " -> ")
			if len(flowSlice) != 2 {
				return labelMap, arr, errors.New("invalid flow: " + flow)
			}

			// split string for source
			srcSlice := strings.Split(flowSlice[0], ":")
			if len(srcSlice) != 2 {
				return labelMap, arr, errors.New("invalid source: " + flowSlice[0])
			}

			// split string for destination
			dstSlice := strings.Split(flowSlice[1], ":")
			if len(dstSlice) != 2 {
				return labelMap, arr, errors.New("invalid destination: " + flowSlice[1])
			}

			// convert ports to integers
			var srcport int
			srcport, err = strconv.Atoi(srcSlice[1])
			if err != nil {
				return
			}
			var dstport int
			dstport, err = strconv.Atoi(dstSlice[1])
			if err != nil {
				return
			}

			// get description
			var (
				dRaw         = description.FindString(l)
				count        int
				dStart, dEnd int
			)

			for i, c := range dRaw {
				if c == ']' {
					count++
				}
				if count == 2 {
					dStart = i + 2
					dEnd = len(dRaw) - 5
					break
				}
			}

			// create alert
			a := &SuricataAlert{
				Timestamp:      utils.TimeToString(t),
				Proto:          nProto,
				SrcIP:          srcSlice[0],
				SrcPort:        srcport,
				DstIP:          dstSlice[0],
				DstPort:        dstport,
				Classification: strings.TrimSuffix(strings.TrimPrefix(classification.FindString(l), "[Classification: "), "]"),
				Description:    dRaw[dStart:dEnd],
			}

			// use attack description instead of classification for labeling
			if useDescription {

				// check if attack class was excluded prior to flipping
				if excluded[a.Classification] {
					continue
				}

				// attack class is not excluded
				// now replace classification with description as requested
				a.Classification = a.Description
			}

			// ensure no alerts with empty classification are collected
			if a.Classification == "" || a.Classification == " " {
				continue
			}

			// count total occurrences of classification
			ClassificationMap[a.Classification]++

			// check if excluded
			if !excluded[a.Classification] {

				// append to collected alerts
				arr = append(arr, a)

				// add to label map
				if _, ok := labelMap[a.Timestamp]; ok {
					// an alert for this timestamp already exists
					// if configured the execution will stop
					// for now the first seen alert for a timestamp will be kept
					duplicates = append(duplicates, a)
				} else {
					labelMap[a.Timestamp] = a
				}
			}
		}
	}

	// if there were duplicates, exit
	if StopOnDuplicateLabels {
		fmt.Println(len(duplicates), "duplicate labels. stopping")

		for _, a := range duplicates {
			tui.Table(os.Stdout, []string{"Field", "Value"}, [][]string{
				[]string{"Timestamp", a.Timestamp},
				[]string{"Proto", a.Proto},
				[]string{"SrcIP", a.SrcIP},
				[]string{"SrcPort", strconv.Itoa(a.SrcPort)},
				[]string{"DstIP", a.DstIP},
				[]string{"DstPort", strconv.Itoa(a.DstPort)},
				[]string{"Classification", a.Classification},
				[]string{"Description", a.Description},
			})
		}

		os.Exit(1)
	}

	fmt.Println(len(duplicates), "alerts ignored in labelMap")

	return labelMap, arr, nil
}
